#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 15 12:13:19 2023

@author: lhackett
"""

import pandas as pd
import geopandas as gpd
from shapely.geometry import Polygon
import numpy as np
import matplotlib.pyplot as plt

meteredproj = 6933
landuse_dir = '/Users/hackettl/Library/CloudStorage/Box-Box/Land use replication/'

#%% Read parcel data, which will give bounds for grid

parcels = gpd.read_file(landuse_dir+"RawData/Pajaro_Parcel_Project.shp")
# project to metered projection for ease
parcels = parcels.to_crs(meteredproj)
# add parcel ID
parcelnums = pd.read_csv(landuse_dir+'out/parcel_IDs_identity.csv')
parcelnums.rename(columns = {'FID_Pajaro' : 'parcelnum'}, inplace = True)
parcels['ID'] = pd.to_numeric(parcels.ID)
parcels = parcels.merge(parcelnums)[['parcelnum', 'geometry']]

# get bounds 
dims = parcels.total_bounds

# clean up
del parcelnums

#%% how large should grid cells be?

areas = parcels[parcels.area*0.000247105 > 2].area # drop parcels  < 2 acres as with land use
totarea = areas.sum()

print(f'Parcel square sides: \n\t median: {areas.median()**.5:.2f} m, \n\t mean: {areas.mean()**.5:.2f} m, \n\t max: {areas.max()**.5:.2f} m')
print(f'Square side if exactly 50: {(totarea/50)**.5:.2f} m')

del areas
del totarea

#%% function to make grid of squares with side length L

def squareGrid(L, dims):
    '''
    Create grid of square polygons with side length L, with extent (dims).
    
    Returns: gpd.GeoDataFrame with geometry and gridID columns
    '''
    xmin, ymin, xmax, ymax = dims
    cols = list(np.arange(xmin, xmax + L, L))
    rows = list(np.arange(ymin, ymax + L, L))
    
    polygons = []
    for x in cols[:-1]:
        for y in rows[:-1]:
            polygons.append(Polygon([(x,y), (x+L, y), (x+L, y+L), (x, y+L)]))
    
    grid = gpd.GeoDataFrame({'geometry':polygons}, crs = parcels.crs)
    print(f'Number of squares: {len(grid.index)}')
    return grid


def mostOverlapped(parc, grids):
    '''
    Function for identifying the largest overlapping grid cell of parcels
    with multiple matches.
    
    Returns: pd.Series with gridID
    '''
    # 1 - isolate grids that overlap the parcel
    subset = grids[np.isin(grids.gridID, parc.gridID)]
    # 2 - get ID of the largest intersection
    idmax = np.argmax(subset.intersection(parc.geometry.iloc[0]).area)
    # Return grid ID
    return subset.iloc[idmax,:]['gridID']

def assignParcels(parcel_df, grid):
    '''
    Function to assign gridIDs to parcels. This involves:
        1 - spatially join the data sets
        2 - for parcels overlying multiple grid cells, return ID with most overlap
    
    Returns: pd.DataFrams with parcelnum, gridID
    '''
    # infer side length
    inferL = str(int(grid.geometry.iloc[0].area**.5))
    # first, spatial join which will result in multiple matches
    parcels_merge = parcel_df.sjoin(grid)

    # find parcels with multiple matches
    nmatch = parcels_merge.groupby('parcelnum').gridID.count()
    nmatch.name = 'nmatch'
    # how many?
    print(f'{len(nmatch[nmatch > 1])} parcels overlying multiple grid cells')
    nmatch = parcels_merge.merge(nmatch.reset_index())
    multmatch = nmatch[nmatch.nmatch > 1]
    
    # assign grid square with largest intersection
    multmatch_grid = multmatch.groupby('parcelnum').apply(lambda x: mostOverlapped(x, grid))
    multmatch_grid.name = 'gridID_'+inferL+'m'
    multmatch_grid = multmatch_grid.reset_index()
    # glue together 1-matches w multmatch
    onematch = nmatch[nmatch.nmatch == 1][['parcelnum', 'gridID']].rename(columns={'gridID' : 'gridID_'+inferL+'m'})
    merged = pd.concat([onematch, multmatch_grid], axis = 0).reset_index(drop=True)
    
    return merged 

#%% test grid creation

fig, ax = plt.subplots()
parcels.plot(ax=ax)
squareGrid(2000, dims).boundary.plot(ax=ax, color = 'red')

del fig
del ax


#%% implement

parcels_assigned = []

for l in [400, 500, 600, 1000]:
    # 1 - create squares
    squares = squareGrid(l, dims).reset_index().rename(columns={'index' : 'gridID'})
    # 2 - assign parcels to grid
    parcels_assigned.append(assignParcels(parcels, squares))
    
# merge
for i in range(len(parcels_assigned)):
    if i == 0:
        dt = parcels_assigned[i]
    else:
        dt = dt.merge(parcels_assigned[i])
        
# clean up
del l
del i
del squares

#%% export

dt.to_csv(landuse_dir+'out/gridIDs.csv', index=False)


    






